#!/usr/bin/env python
# coding=utf-8
# 对EVM的结果进行从新命名

from BCBio import GFF
import argparse
import sys

parser = argparse.ArgumentParser(
    description='''对EVM的结果进行从新命名，
    用法:
    change_evm_id.py -i evm.gff3 -p NIG_PB -o genes.gff3 -n 1
    由大天才于2021年7月6日创建于浙江农业大学''')

parser.add_argument('-i',
                help='必须给定，输入的gff文件 需是evm的输出')

parser.add_argument('-p',
                help='必须给定，确定基因组版本的前缀')

parser.add_argument('-o',
                help='必须给定，输出文件')

parser.add_argument('-n',
                help='起始编号，默认为1')



args = parser.parse_args()

if not args.i or not args.o:
    parser.print_help()
    sys.exit()



infile = args.i


if not args.p:
    prefix='EVM'
else:
    prefix = args.p


outfile = args.o




inhandle = open(infile,'r',encoding='utf-8')
outhandle = open(outfile,'w',encoding='utf-8')



if not args.n:
    n =1
else:
    n = int(args.n)





# 读取 内容

lista = []
for rec in GFF.parse(inhandle):


    feature_lista = []

    for feature in rec.features:

        if feature.type == 'gene':

            feature.id = prefix + '.gene' +str(n)

            feature.qualifiers['ID'] = feature.id
            feature.qualifiers['Name'] = feature.id

            n+=1

            mRNA_index = 1

            mRNA_lista = []

            for mRNA in feature.sub_features:

                if mRNA.type == 'mRNA':

                    mRNA.id = feature.id+'.mRNA'+str(mRNA_index)

                    mRNA.qualifiers['ID'] = mRNA.id
                    mRNA.qualifiers['Name'] = mRNA.id
                    
                    mRNA.qualifiers['Parent'] = feature.id

                    mRNA_index += 1

                    exonOrCDS_lista = []

                    exon_index = 1
                    cds_index = 1

                    for exonOrCDS in mRNA.sub_features:
                        if exonOrCDS.type == 'CDS':

                            exonOrCDS.id = mRNA.id +'.CDS'+ str(cds_index)

                            cds_index += 1

                        elif exonOrCDS.type == 'exon':

                            exonOrCDS.id = mRNA.id +'.exon'+ str(cds_index)

                            exon_index += 1

                        else:
                            continue

                        exonOrCDS.qualifiers['ID'] = exonOrCDS.id

                        exonOrCDS.qualifiers['Name'] = exonOrCDS.id
                
                        exonOrCDS.qualifiers['Parent'] = mRNA.id

                        exonOrCDS_lista.append(exonOrCDS)

                    mRNA.sub_features = exonOrCDS_lista

                    mRNA_lista.append(mRNA)

            feature.sub_features = mRNA_lista

            feature_lista.append(feature)
            
    rec.features = feature_lista

    lista.append(rec)

GFF.write(lista,outhandle)

       
    


